#include "Environment.hh"
#include <iostream>

using namespace std;

void StackFrame::bindDecl(Decl* decl, int64_t val) {
	mVars[decl] = val;
}

int64_t StackFrame::getDeclVal(Decl * decl) {
	auto dit = mVars.find(decl);
	if (dit == mVars.end()) {
		VarDecl * vardecl = dyn_cast<VarDecl>(decl);
		if (vardecl != nullptr)
			throw logic_error("Variable " + vardecl->getNameAsString()
				+ " not bound to any value");
		else
			throw logic_error("Declaration not bound to any value");
	}
	return mVars.find(decl)->second;
}

void StackFrame::bindStmt(Stmt * stmt, int64_t val) {
	mExprs[stmt] = val;
}

int64_t StackFrame::getStmtVal(Stmt * stmt) {
	auto eit = mExprs.find(stmt);
	if (eit == mExprs.end()) {
		throw logic_error("Statement not bound to any value");
	}
	return mExprs[stmt];
}

void StackFrame::bindPtr(Stmt * stmt, int64_t addr) {
	mPtrs[stmt] = addr;
}

int64_t StackFrame::getPtr(Stmt * stmt) {
	auto pit = mPtrs.find(stmt);
	if (pit == mPtrs.end()) {
		throw logic_error("Statement not bound to any address");
	}
	return mPtrs[stmt];
}

void StackFrame::setPC(Stmt * stmt) {
	mPC = stmt;
}

Stmt * StackFrame::getPC() {
	return mPC;
}

Heap::~Heap() {
	for (auto & h : heap)
		delete [] (int8_t*)(h.first);
}

/*
 * Allocate memory on the heap and initialize to zero.
 */
void * Heap::Malloc(int64_t size) {
	int64_t i;
	if (size <= 0)
		throw runtime_error("Invalid memory allocation size");
	void * p = new int8_t[size];
	heap[p] = size;
	for (i = 0; i < size; ++i)
		((int8_t*)p)[i] = 0;
	return p;
}

void Heap::Free(void * p) {
	auto it = heap.find(p);
	if (it == heap.end())
		throw runtime_error("Invalid free operation");
	delete [] (int8_t*)(it->first);
	heap.erase(it);
}

/// Initialize the Environment
void Environment::init(TranslationUnitDecl * unit) {
	// Create the global "stack frame"
	mStack.push_back(StackFrame());
}

FunctionDecl * Environment::getEntry() {
	return mEntry;
}

void Environment::binop(BinaryOperator *bop) {
	Expr * left = bop->getLHS();
	Expr * right = bop->getRHS();
	BinaryOperator::Opcode opcode = bop->getOpcode();
	int64_t val = 0, rval = mStack.back().getStmtVal(right), lval;

	if (left->getType()->isPointerType()
		&& (right->getType()->isIntegerType()
			|| right->getType()->isCharType())) {
		QualType lpetype = left->getType()->getPointeeType();
		if (lpetype->isIntegerType())
			rval *= sizeof(int);
		else if (lpetype->isCharType())
			rval *= sizeof(char);
		else if (lpetype->isPointerType())
			rval *= sizeof(void*);
		else
			throw logic_error("Unsupported pointer element type"
				" in binary operation");
	}

	if (bop->isAssignmentOp()) {
		if (opcode != BO_Assign)
			// Operators such as `+=` are not supported yet
			throw logic_error("Only simple assignment operation supported yet");
		val = rval;
		mStack.back().bindStmt(left, val);
		if (DeclRefExpr * declexpr = dyn_cast<DeclRefExpr>(left)) {
			Decl * decl = declexpr->getFoundDecl();
			// Bind value to variable declaration (global)
			bindDecl(decl, val);
		} else if (ArraySubscriptExpr * asexpr = dyn_cast<ArraySubscriptExpr>(left)) {
			int64_t addr = mStack.back().getPtr(asexpr);
			QualType atype = asexpr->getType();
			if (atype->isCharType()) {
				*((char*)(addr)) = (char)(val);
			} else if (atype->isIntegerType()) {
				*((int*)(addr)) = (int)(val);
			} else if (atype->isPointerType()) {
				*((void**)(addr)) = (void*)(val);
			} else {
				throw logic_error("Unsupported array element type in assignment");
			}
		} else if (UnaryOperator * uop = dyn_cast<UnaryOperator>(left)) {
			if (uop->getOpcode() == UO_Deref) {
				int64_t addr = mStack.back().getPtr(uop);
				QualType uotype = uop->getType();
				if (uotype->isIntegerType()) {
					*((int*)(addr)) = (int)(val);
				} else if (uotype->isCharType()) {
					*((char*)(addr)) = (char)(val);
				} else if (uotype->isPointerType()) {
					*((void**)(addr)) = (void*)(val);
				} else {
					throw logic_error("Unsupported dereference type in assignment");
				}
			} else {
				throw logic_error("Unsupported left-hand side in assignment");
			}
		}
		else
			throw logic_error("Unsupported left-hand side in assignment");
	} else {
		// support other binary operations: +, -, *, /, <, >, ==
		lval = mStack.back().getStmtVal(left);
		switch (opcode) {
		case BO_Add: val = lval + rval; break;
		case BO_Sub: val = lval - rval; break;
		case BO_Mul: val = lval * rval; break;
		case BO_Div:
			if (rval == 0)
				throw runtime_error("Division by zero");
			val = lval / rval;
			break;
		case BO_LT: val = (lval < rval) ? 1 : 0; break;
		case BO_GT: val = (lval > rval) ? 1 : 0; break;
		case BO_GE: val = (lval >= rval) ? 1 : 0; break;
		case BO_LE: val = (lval <= rval) ? 1 : 0; break;
		case BO_EQ: val = (lval == rval) ? 1 : 0; break;
		default: 
			throw logic_error("Unsupported binary operation");
		}
	}
	mStack.back().bindStmt(bop, val);
}

void Environment::unop(UnaryOperator * uop) {
	Expr * subexpr = uop->getSubExpr();
	int64_t sval = mStack.back().getStmtVal(subexpr);
	int64_t val = 0;
	switch (uop->getOpcode()) {
	case UO_Minus:
		val = -sval;
		break;
	case UO_Plus:
		val = +sval;
		break;
	case UO_Deref: {
		QualType uotype = uop->getType();
		if (uotype->isIntegerType()) {
			val = *((int*)(sval));
		} else if (uotype->isCharType()) {
			val = *((char*)(sval));
		} else if (uotype->isPointerType()) {
			val = int64_t(*((void**)(sval)));
		} else {
			throw logic_error("Unsupported dereference type");
		}
		mStack.back().bindPtr(uop, sval);
		break;
	}
	default:
		throw logic_error("Unsupported unary operation");
	}
	mStack.back().bindStmt(uop, val);
}

void Environment::uett(UnaryExprOrTypeTraitExpr * uett) {
	mStack.back().setPC(uett);
	if (uett->getKind() == UETT_SizeOf) {
		QualType uettype = uett->getTypeOfArgument();
		CharUnits size = context.getTypeSizeInChars(uettype);
		mStack.back().bindStmt(uett, size.getQuantity());
	} else
		throw logic_error("Unsupported UETT kind");
}

void Environment::decl(DeclStmt * declstmt) {
	for (DeclStmt::decl_iterator it = declstmt->decl_begin(),
			ie = declstmt->decl_end();
			it != ie; ++ it) {
		Decl * decl = *it;
		// Bind variable declaration
		if (VarDecl * vdecl = dyn_cast<VarDecl>(decl)) {
			vardecl(vdecl);
		}
		// Function declaration
		else if (FunctionDecl * fdecl = dyn_cast<FunctionDecl>(decl)) {
			// Handle function declaration
			this->fdecl(fdecl);
		}
	}
}

void Environment::vardecl(VarDecl * vdecl) {
	if (vdecl->getType()->isIntegerType()
		|| vdecl->getType()->isPointerType()
		|| vdecl->getType()->isCharType()) {
		Expr * init = vdecl->getInit();
		if (init != nullptr) {
			int64_t val = mStack.back().getStmtVal(init);
			mStack.back().bindDecl(vdecl, val);
		} else {
			mStack.back().bindDecl(vdecl, 0);
		}
	} else if (vdecl->getType()->isArrayType()) {
		void * aaddr = nullptr;
		auto array = dyn_cast<ConstantArrayType>(
			vdecl->getType().getTypePtr());
		QualType elemtype = array->getElementType();
		int64_t asize = array->getSize().getSExtValue();
		if (elemtype->isCharType())
			aaddr = mHeap.Malloc(asize * sizeof(char));
		else if (elemtype->isIntegerType())
			aaddr = mHeap.Malloc(asize * sizeof(int));
		else if (elemtype->isPointerType())
			aaddr = mHeap.Malloc(asize * sizeof(void*));
		else
			throw logic_error("Unsupported array element type");
		mStack.back().bindDecl(vdecl, (int64_t)(aaddr));
	}
}

void Environment::fdecl(FunctionDecl * fdecl) {
	const char * fname = fdecl->getName().data();
	if (fname == nullptr)
		/*
		 * Note: We suppose that every `FunctionDecl` has
		 * a valid name in this interpreter.
		 */
		throw std::logic_error("Invalid FunctionDecl name");
	if (strcmp(fname, "FREE") == 0) mFree = fdecl;
	else if (strcmp(fname, "MALLOC") == 0) mMalloc = fdecl;
	else if (strcmp(fname, "GET") == 0) mInput = fdecl;
	else if (strcmp(fname, "PRINT") == 0) mOutput = fdecl;
	else if (strcmp(fname, "main") == 0) mEntry = fdecl;
}

void Environment::declref(DeclRefExpr * declref) {
	mStack.back().setPC(declref);
	// Check if the referenced declaration is a variable
	if (declref->getType()->isIntegerType()
		|| declref->getType()->isCharType()
		|| declref->getType()->isArrayType()
		|| declref->getType()->isPointerType()) {
		Decl* decl = declref->getFoundDecl();
		VarDecl * vardecl = dyn_cast<VarDecl>(decl);
		int64_t val = 0;
		bool flag;
		size_t stki = mStack.size();
		string lex_msg;
		do {
			--stki;
			flag = true;
			try {
				val = mStack[stki].getDeclVal(decl);
			} catch (logic_error & lex) {
				lex_msg = lex.what();
				flag = false;	// failed
			}
			/*
			 * For variable reference, we search from the top stack frame
			 * to the bottom stack frame to find the variable binding.
			 */
		} while (vardecl != nullptr && stki > 0u && flag == false);
		if (flag == false)
			throw logic_error(lex_msg);
		mStack.back().bindStmt(declref, val);
	} else if (declref->getType()->isFunctionType()) {
		;
	} else {
		throw logic_error("Unsupported DeclRefExpr type");
	}
}

/*
 * For assignment to variable declaration, we search from the top stack frame.
 * If found, bind the value to the variable declaration in that stack frame.
 * If not found, throw an exception.
 */
void Environment::bindDecl(Decl* decl, int64_t val) {
	string lex_msg;
	for (auto sit = mStack.rbegin(); sit != mStack.rend(); ++sit) {
		try {
			sit->getDeclVal(decl);
			sit->bindDecl(decl, val);
			return;
		} catch (logic_error & lex) {
			// Not found in this stack frame, continue searching
			lex_msg = lex.what();
		}
	}
	throw logic_error(lex_msg);
}

void Environment::cast(CastExpr * castexpr) {
	mStack.back().setPC(castexpr);
	if ((castexpr->getType()->isIntegerType()
		|| castexpr->getType()->isCharType()
		|| castexpr->getType()->isPointerType())
		&& !castexpr->getType()->isFunctionPointerType()) {
		Expr * expr = castexpr->getSubExpr();
		int64_t val = mStack.back().getStmtVal(expr);
		mStack.back().bindStmt(castexpr, val );
	} else if (castexpr->getType()->isFunctionType()
		|| castexpr->getType()->isFunctionPointerType()) {
		;
	} else {
		throw logic_error("Unsupported CastExpr type");
	}
}

/*
 * Handle function call.
 * If the called function is one of the built-in functions,
 * handle it here and return 0.
 * Otherwise, return a non-zero value to indicate that
 * the function call needs to be handled by the caller.
 */
long Environment::call(CallExpr * callexpr) {
	mStack.back().setPC(callexpr);
	int64_t val = 0;
	long rt = 0l;
	FunctionDecl * callee = callexpr->getDirectCallee();
	if (callee->isDefined())
		callee = callee->getDefinition();
	if (callee == mInput) {
		llvm::outs() << "Please Input an Integer Value : ";
		std::cin >> val;

		mStack.back().bindStmt(callexpr, val);
	} else if (callee == mOutput) {
		Expr * decl = callexpr->getArg(0);
		val = mStack.back().getStmtVal(decl);
		llvm::outs() << val;
	} else if (callee == mMalloc) {
		Expr * decl = callexpr->getArg(0);
		val = mStack.back().getStmtVal(decl);
		val = int64_t(mHeap.Malloc(val));
		mStack.back().bindStmt(callexpr, val);
	} else if (callee == mFree) {
		Expr * decl = callexpr->getArg(0);
		val = mStack.back().getStmtVal(decl);
		mHeap.Free((void*)(val));
	} else {
		unsigned int na = callexpr->getNumArgs(), ia;
		if (na != callee->getNumParams())
			throw runtime_error("Number of arguments does not match");
		StackFrame calleesf;
		for (ia = 0u; ia < na; ++ia)
			calleesf.bindDecl(
				callee->getParamDecl(ia),
				mStack.back().getStmtVal(callexpr->getArg(ia))
			);
		mStack.push_back(calleesf);	// for the callee
		rt = 9;
	}
	return rt;
}

void Environment::arraysub(ArraySubscriptExpr * asexpr) {
	mStack.back().setPC(asexpr);
	Expr * baseexpr = asexpr->getBase(),
		* indexexpr = asexpr->getIdx();
	int64_t baseaddr = mStack.back().getStmtVal(baseexpr),
		index = mStack.back().getStmtVal(indexexpr);
	QualType atype = asexpr->getType();
	int64_t val = 0l, addr = 0l;
	if (atype->isCharType()) {
		val = ((char*)(baseaddr))[index];
		addr = baseaddr + index * int64_t(sizeof(char));
	} else if (atype->isIntegerType()) {
		val = ((int*)(baseaddr))[index];
		addr = baseaddr + index * int64_t(sizeof(int));
	} else if (atype->isPointerType()) {
		val = int64_t(((void**)(baseaddr))[index]);
		addr = baseaddr + index * int64_t(sizeof(void*));
	} else {
		throw logic_error("Unsupported array element type");
	}
	mStack.back().bindStmt(asexpr, val);
	mStack.back().bindPtr(asexpr, addr);
}

void Environment::retstmt(ReturnStmt * retstmt) {
	mStack.back().setPC(retstmt);
	Expr * retexpr = retstmt->getRetValue();
	int64_t val = 0;
	if (retexpr != nullptr
		&& retexpr->getType()->isVoidType() == false) {
		val = mStack.back().getStmtVal(retexpr);
		mStack.back().bindRetVal(val);
	}
}

void Environment::freturn(CallExpr * callexpr) {
	int64_t retv = mStack.back().getRetVal();
	mStack.pop_back();	// pop the callee stack frame
	if (!callexpr->getDirectCallee()->getReturnType()->isVoidType())
		mStack.back().bindStmt(callexpr, retv);
}

long Environment::cond(Expr * condexpr) {
	return mStack.back().getStmtVal(condexpr);
}

void Environment::intlit(IntegerLiteral * intlit) {
	mStack.back().setPC(intlit);
	int val = intlit->getValue().getSExtValue();
	mStack.back().bindStmt(intlit, val);
}

void Environment::charlit(CharacterLiteral * charlit) {
	mStack.back().setPC(charlit);
	int val = charlit->getValue();
	mStack.back().bindStmt(charlit, val);
}

void Environment::paren(ParenExpr * parenexpr) {
	mStack.back().setPC(parenexpr);
	Expr * subexpr = parenexpr->getSubExpr();
	int64_t val = mStack.back().getStmtVal(subexpr);
	mStack.back().bindStmt(parenexpr, val);
}
